package com.zzh.demo.agentDemo.transformer;

import com.zzh.demo.agentDemo.utils.ConfigUtils;
import com.zzh.demo.agentDemo.visitor.TimeStatisticsVisitor;
import jdk.internal.org.objectweb.asm.ClassReader;
import jdk.internal.org.objectweb.asm.ClassVisitor;
import jdk.internal.org.objectweb.asm.ClassWriter;
import jdk.internal.org.objectweb.asm.Opcodes;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Set;

/**
 * description： <br>
 * createTime: 2024/2/29 10:38 <br>
 *
 * @author zzh
 */
public class ClassMethodRunTimeTransformer implements ClassFileTransformer {
    static Set<String> agentPackagePrefixes = new HashSet(Arrays.asList(
            ConfigUtils.getStrProperties("agentPackagePrefix", "com.asiainfo").split(",")));
    static {
        System.out.println("agentPackagePrefixes = " + agentPackagePrefixes);
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        if (check(className)) {
//            System.out.println(className);
            // 使用ASM框架进行字节码转换
            ClassReader cr = new ClassReader(classfileBuffer);
            ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_FRAMES);
            ClassVisitor cv = new TimeStatisticsVisitor(Opcodes.ASM5, cw);
            cr.accept(cv, ClassReader.SKIP_FRAMES | ClassReader.SKIP_DEBUG);
            return cw.toByteArray();
        }
        return classfileBuffer;
    }

    private boolean check(String className) {
        for (String agentPackagePrefix : agentPackagePrefixes) {
            if(className.startsWith(agentPackagePrefix)) {
                return true;
            }
        }
        return false;
    }
}
